package com.hxx.common.core.wcontour.util;

import com.hxx.common.core.wcontour.base.TrainData;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.util.Iterator;
import java.util.Map;

/**
 * 热力图
 *
 * @author hxd
 */
public class GeoHeatUtil {

    /**
     * 网格数据生成热力图（输出到img模版里）
     *
     * @param trainData  网格数据
     * @param levelProps 色阶,结构如：{0.1:"#a5f38d"}，需要是LinkedHashMap，且key值是递增（升序）
     * @param tempPath   图片模版路径
     * @param outPath    输出地址
     * @throws IOException
     */
    public static void gridToPng(TrainData trainData, Map<Double, String> levelProps, String tempPath,
                                     String outPath) throws IOException {
        // 网格数据
        double[][] data = trainData.getData();
        // 这里我默认的是广东省的起点经纬度，自定义可以从shp获取或者入参数
        double minX = 109.65, minY = 20.21;
        /**
         * 网格和图片模版对应的起点坐标，由于1像素=0.01经纬度，所以坐标差获取网格起点下标
         * 需要注意：这里的网格数据完全包含img模板范围，如果不固定范围根据业务自行调整
         */
        int startX = (int) ((minX - trainData.get_X()[0]) / 0.01);
        int startY = (int) ((minY - trainData.get_Y()[0]) / 0.01);

        double[] levelPropsArr = new double[levelProps.size()];
        Iterator iter = levelProps.entrySet().iterator();
        int index = 0;
        while (iter.hasNext()) {
            Map.Entry entry = (Map.Entry) iter.next();
            levelPropsArr[index++] = (double) entry.getKey();
        }
        // 默认取第一个颜色
        String colorHex = levelProps.get(levelPropsArr[0]);

        File file = new File(tempPath);
        BufferedImage bi = ImageIO.read(file);
        int width = bi.getWidth(), height = bi.getHeight();
        for (int h = 0; h < height; h++) {
            for (int w = 0; w < width; w++) {
                // 转16进制
                int pixel = bi.getRGB(w, h);
                // 判断RGB是否为(0,0,255)，然后进行像素颜色替换
                if ((pixel & 0xff0000) >> 16 == 0 && (pixel & 0xff00) >> 8 == 0 && (pixel & 0xff) == 255) {
                    // 当前坐标值
                    double value = data[startY + h][startX + w];
                    for (int i = levelPropsArr.length - 1; i >= 0; i--) {
                        if (value >= levelPropsArr[i]) {
                            // 符合条件
                            colorHex = levelProps.get(levelPropsArr[i]);
                            break;
                        }
                    }
                    int newPixel = ColorUtil.hexToInt(colorHex);
                    bi.setRGB(w, h, newPixel);
                }
            }
        }
        ImageIO.write(bi, "png", new File(outPath));
    }

    /**
     * 网格数据生成热力图
     * @param trainData 网格数据
     * @param levelProps 色阶,结构如：{0.1:"#a5f38d"}，需要是LinkedHashMap，且key值是递增（升序）
     * @param outPath 图片生成地址
     * @throws IOException
     */
    public static void gridToPng(TrainData trainData, Map<Double, String> levelProps, String outPath) throws IOException {
        // 网格数据
        double[][] data = trainData.getData();
        int width = trainData.getSize()[0], height = trainData.getSize()[1];

        double[] levelPropsArr = new double[levelProps.size()];
        Iterator iter = levelProps.entrySet().iterator();
        int index = 0;
        while (iter.hasNext()) {
            Map.Entry entry = (Map.Entry) iter.next();
            levelPropsArr[index++] = (double) entry.getKey();
        }
        // 默认取第一个颜色
        String colorHex = levelProps.get(levelPropsArr[0]);

        BufferedImage bi = new BufferedImage(width, height,BufferedImage.TYPE_INT_RGB);
        for (int h = 0; h < height; h++) {
            for (int w = 0; w < width; w++) {
                // 转16进制
                double value = data[h][w];
                for (int i = levelPropsArr.length - 1; i >= 0; i--) {
                    if (value >= levelPropsArr[i]) {
                        // 符合条件
                        colorHex = levelProps.get(levelPropsArr[i]);
                        break;
                    }
                }
                int newPixel = ColorUtil.hexToInt(colorHex);
                bi.setRGB(w, h, newPixel);
            }
        }
        ImageIO.write(bi, "png", new File(outPath));
    }
}
